from torch import nn
import torch
from torchvision.models.vgg import vgg11_bn


class DDPM(nn.Module):
    def __init__(self, out_dim=3072):
        super().__init__()
        self.conv = nn.Sequential(
            vgg11_bn().features,
            nn.Flatten(),
        )
        self.step = nn.Sequential(
            nn.Linear(1, 32)
        )

        self.reg = nn.Sequential(
            nn.Linear(512 + 32, 4096, bias=True),
            nn.BatchNorm1d(4096),
            nn.ReLU(),
            nn.Linear(4096, 4096),
            nn.BatchNorm1d(4096),
            nn.Linear(4096, out_dim),
            nn.Sigmoid()
        )


    def forward(self, x, t):
        img = torch.cat((self.conv(x), self.step(t)), dim=-1)
        return self.reg(img) 
        